#plotSpread.py

import Quandl
import matplotlib.pyplot as plt
import time
import datetime
import os.path
import pandas as pd
import numpy as np
from pandas import DataFrame
from sklearn import preprocessing

token = "piaztMKBTwk8kVkxAvS9"
monthCodes = {'F' : 1, 'G' : 2, 'H' : 3, 'J' : 4, 'K' : 5, 'M' : 6, 'N' : 7, 'Q' : 8, 'U' : 9, 'V' : 10, 'X' : 11, 'Z' : 12}

def getSymbol(name):
    if os.path.isfile('./data/' + name[4:] + '.csv'):
        print 'from file ... '
        symbol = pd.read_csv('./data/' + name[4:] + '.csv', index_col=0)
        symbol.index = pd.to_datetime(symbol.index)
    else:
        print 'from Quandl ...'
        symbol = Quandl.get(name, authtoken=token)
        symbol.to_csv('./data/' + name[4:] + '.csv')

    return symbol

def getSpread(longSymbol, shortSymbol):
    # load data
    print 'Loading ' + longSymbol + ' and ' + shortSymbol + ' ...'
    longLeg = getSymbol(longSymbol)
    shortLeg = getSymbol(shortSymbol)

    # calculate spread
    spread = longLeg['Settle'].subtract(shortLeg['Settle'])

    # get rid of nan
    spread = spread.dropna()

    return spread

def getSpreadStat(longSymbol, shortSymbol, firstYear, lastYear):

    df = pd.DataFrame()
    lastValidDate = datetime.datetime(1975,1,1)
    min_max_scaler = preprocessing.MinMaxScaler()

    for year in range(firstYear, lastYear):
        symbolShortLeg = shortSymbol + str(year)

        # check if year wrap is needed
        if monthCodes[longSymbol[-1]] < monthCodes[shortSymbol[-1]]:
            symbolLongLeg = longSymbol + str(year+1)
        else:
            symbolLongLeg = longSymbol + str(year)

        print symbolLongLeg + ' ' + symbolShortLeg

        spread = getSpread(symbolLongLeg, symbolShortLeg)
        spread = spread.resample("D", fill_method='bfill')

        if lastValidDate.day < spread.last_valid_index().day:
            lastValidDate = spread.last_valid_index()

        df_tmp = DataFrame(spread.get_values(), index = spread.index, columns = ['Settle'])
        df_tmp['scale'] = min_max_scaler.fit_transform(df_tmp['Settle'].get_values())

        df = df.append(df_tmp[datetime.datetime(year-1, monthCodes[shortSymbol[-1]], 28):])

    df['key'] = [x.strftime('%m-%d') for x in df.index]
    df = df.groupby('key').mean()

    df_tmp = df[lastValidDate.strftime('%m-%d'):]
    df_tmp = df_tmp.append(df[:lastValidDate.strftime('%m-%d')])

    df_tmp.to_csv('spread.csv')
    return df_tmp


def getSpreadStatLastYears(longSymbol, shortSymbol, years = 15):

    timeNow = datetime.date.fromtimestamp(time.time())
    return timeNow

#for year in years:
 #   symbolLongLeg = "CME/CK" + str(year)
 #   symbolShortLeg = "CME/CZ" + str(year-1)
    
    #spread = getSpread(symbolLongLeg, symbolShortLeg)
    #spread.plot()
    #plt.show()

# spread0 = getSpread("CME/CK2012", "CME/CZ2011")
# print spread0
# spread0.plot()
# plt.show()

#spread1 = getSpread("CME/CK2013", "CME/CZ2013")
#print '...'
# test = getSpreadStat("CME/CZ", "CME/CU", 1985, 2005)
# tmp = test.groupby('key').mean()
# tmp = tmp.shift(90)
# print tmp
#
# tmp.plot(y = 'scale')
# plt.title("CME/CK CME/CZ")
# plt.show()
# test = getSpreadStat("CME/CK", "CME/CZ", 1999, 2014)
# test.groupby('key').mean().plot(y = 'scale')
# plt.title("CME/CU CME/CZ")
# plt.show()
